import torch
import torch.nn.functional as F
from mmengine.model import BaseModel

from torchvision.models.segmentation import deeplabv3_resnet50


class MMDeeplabV3(BaseModel):
    def __init__(self, num_classes):
        super().__init__()
        self.deeplab = deeplabv3_resnet50()
        self.deeplab.classifier[4] = torch.nn.Conv2d(
            256, num_classes, kernel_size=(1, 1), stride=(1, 1)
        )

    def forward(self, imgs, data_samples=None, mode="tensor"):
        x = self.deeplab(imgs)["out"]
        if mode == "loss":
            return {"loss": F.cross_entropy(x, data_samples["labels"])}
        elif mode == "predict":
            return x, data_samples
